import asyncio
import os
from openai import OpenAI
from dotenv import load_dotenv
from contextlib import AsyncExitStack
import json
from fastmcp import Client

"""
AI 调用 Streamable HTTP 模式的 MCP Sever 上的 tools

运行步骤：
1. 在根目录执行 `python server.py`
2. 在 test 目录执行 `python agent_http.py `
"""


# 系统提示词
system_prompt = """
In this environment you have access to a set of tools you can use to answer the user's question. You can use one or more tools per message, and will receive the result of that tool use in the user's response. You use tools step-by-step to accomplish a given task, with each tool use informed by the result of the previous tool use.

## Tool Use Formatting

Tool use is formatted using XML-style tags. The tool name is enclosed in opening and closing tags, and each parameter is similarly enclosed within its own set of tags. Here's the structure:

<tool_use>
  <name>{tool_name}</name>
  <arguments>{json_arguments}</arguments>
</tool_use>

The tool name should be the exact name of the tool you are using, and the arguments should be a JSON object containing the parameters required by that tool. For example:
<tool_use>
  <name>python_interpreter</name>
  <arguments>{\"code\": \"5 + 3 + 1294.678\"}</arguments>
</tool_use>

The user will respond with the result of the tool use, which should be formatted as follows:

<tool_use_result>
  <name>{tool_name}</name>
  <result>{result}</result>
</tool_use_result>

The result should be a string, which can represent a file or any other output type. You can use this result as input for the next action.
For example, if the result of the tool use is an image file, you can use it in the next action like this:

<tool_use>
  <name>image_transformer</name>
  <arguments>{\"image\": \"image_1.jpg\"}</arguments>
</tool_use>

Always adhere to this format for the tool use to ensure proper parsing and execution.

## Tool Use Examples

Here are a few examples using notional tools:
---
User: Generate an image of the oldest person in this document.

Assistant: I can use the document_qa tool to find out who the oldest person is in the document.
<tool_use>
  <name>document_qa</name>
  <arguments>{\"document\": \"document.pdf\", \"question\": \"Who is the oldest person mentioned?\"}</arguments>
</tool_use>

User: <tool_use_result>
  <name>document_qa</name>
  <result>John Doe, a 55 year old lumberjack living in Newfoundland.</result>
</tool_use_result>

Assistant: I can use the image_generator tool to create a portrait of John Doe.
<tool_use>
  <name>image_generator</name>
  <arguments>{\"prompt\": \"A portrait of John Doe, a 55-year-old man living in Canada.\"}</arguments>
</tool_use>

User: <tool_use_result>
  <name>image_generator</name>
  <result>image.png</result>
</tool_use_result>

Assistant: the image is generated as image.png

---
User: \"What is the result of the following operation: 5 + 3 + 1294.678?\"

Assistant: I can use the python_interpreter tool to calculate the result of the operation.
<tool_use>
  <name>python_interpreter</name>
  <arguments>{\"code\": \"5 + 3 + 1294.678\"}</arguments>
</tool_use>

User: <tool_use_result>
  <name>python_interpreter</name>
  <result>1302.678</result>
</tool_use_result>

Assistant: The result of the operation is 1302.678.

---
User: \"Which city has the highest population , Guangzhou or Shanghai?\"

Assistant: I can use the search tool to find the population of Guangzhou.
<tool_use>
  <name>search</name>
  <arguments>{\"query\": \"Population Guangzhou\"}</arguments>
</tool_use>

User: <tool_use_result>
  <name>search</name>
  <result>Guangzhou has a population of 15 million inhabitants as of 2021.</result>
</tool_use_result>

Assistant: I can use the search tool to find the population of Shanghai.
<tool_use>
  <name>search</name>
  <arguments>{\"query\": \"Population Shanghai\"}</arguments>
</tool_use>

User: <tool_use_result>
  <name>search</name>
  <result>26 million (2019)</result>
</tool_use_result>
Assistant: The population of Shanghai is 26 million, while Guangzhou has a population of 15 million. Therefore, Shanghai has the highest population.


## Tool Use Available Tools
Above example were using notional tools that might not exist for you. You only have access to these tools:
<tools>
{available_tools}
</tools>

## Tool Use Rules
Here are the rules you should always follow to solve your task:
1. Always use the right arguments for the tools. Never use variable names as the action arguments, use the value instead.
2. Call a tool only when needed: do not call the search agent if you do not need information, try to solve the task yourself.
3. If no tool call is needed, just answer the question directly.
4. Never re-do a tool call that you previously did with the exact same parameters.
5. For tool use, MARK SURE use XML tag format as shown in the examples above. Do not use any other format.

# User Instructions

Response in user query language.
Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000.
"""


# 加载 .env 文件
load_dotenv()

# 使用 config 配置方式
config = {
    "mcpServers": {
        "server": {
            "url": "http://127.0.0.1:3002/mcp_atlas", # 服务器地址，需要根据 MCP Server 配置修改 
            "transport": "streamable-http"
        }
    }
}

class MCPClient:
    def __init__(self):
        """初始化 MCP 客户端"""
        self.exit_stack = AsyncExitStack()
        self.api_key = os.getenv("API_KEY")  # 读取 OpenAI API Key
        self.base_url = os.getenv("BASE_URL")  # 读取 BASE URL
        self.model = os.getenv("MODEL")  # 读取 model

        if not self.base_url: 
            raise ValueError("未找到 API BASE_URL. 请在 .env 文件中配置 BASE_URL")

        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)

    async def process_query(self, query: str) -> str:
        """
        调用大模型处理用户查询并根据返回的 tools 列表调用对应工具
        """
        # 获取工具列表 - JSON 格式
        available_tools = [{
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "input_schema": tool.inputSchema
            }
        } for tool in self.tools]
        # 获取工具列表 - XML 格式（用在 系统提示词 中）
        available_tools_xml = [
          f"\n  <tool>\n    <name>{tool.name}</name>\n    <description>{tool.description}</description>\n    <arguments>{tool.inputSchema}</arguments>\n  </tool>\n"
        for tool in self.tools]

        # 待发给 AI 的消息
        messages = [
            {"role": "system", "content": system_prompt.replace("{available_tools}", "".join(available_tools_xml))}, 
            {"role": "user", "content": query}
        ]

        print(f"debug info === MCP server 返回的工具列表 available_tools：{available_tools}")
        print(f"debug info === 拼接用户输入后的 messages：{messages}")

        # 请求 OpenAI 模型处理
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            tools=available_tools
        )

        # 处理返回的内容
        content = response.choices[0]
        print(f"debug info === AI 返回的 content： {content}")
        if content.finish_reason == "tool_calls":
            # 执行工具调用
            tool_call = content.message.tool_calls[0]
            tool_name = tool_call.function.name
            tool_args = json.loads(tool_call.function.arguments)

            # 执行工具
            result = await self.session.call_tool(tool_name, tool_args)
            print(f"\n\n[Calling tool {tool_name} with args {tool_args}]\n\n")
            print(f"debug info === MCP tools 调用结果 result：{result}")

            # 将模型返回的原始消息和工具执行的结果都添加到 messages 中
            messages.append(content.message.model_dump())
            messages.append({
                "role": "tool",
                "content": result.content[0].text,
                "tool_call_id": tool_call.id,
            })

            # 将上面的结果再返回给大模型生产最终的结果
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
            )
            return response.choices[0].message.content

        return content.message.content

    async def chat_loop(self):
        """运行交互式聊天循环"""
        print("MCP 客户端已启动！输入 'exit' 退出")

        while True:
            try:
                query = input("问: ").strip()
                if query.lower() == 'exit':
                    break

                response = await self.process_query(query)
                print(f"AI回复: {response}")

            except Exception as e:
                print(f"发生错误: {str(e)}")

    async def clean(self):
        """清理资源"""
        await self.exit_stack.aclose()

    async def connect_to_server(self, server_script_path: str):
        """连接到 MCP 服务器"""
        self.session = await self.exit_stack.enter_async_context(Client(config))

    async def list_tools(self):
        """列出所有工具"""
        # 列出 MCP 服务器上的工具
        response = await self.session.list_tools()
        self.tools = response


async def main():
    # 启动并初始化 MCP 客户端
    client = MCPClient()
    try:
        # 连接到 MCP 服务器
        await client.connect_to_server('server.py')
        # 列出 MCP 服务器上的工具
        await client.list_tools()
        # 运行交互式聊天循环，处理用户对话
        await client.chat_loop()
    finally:
        # 清理资源
        await client.clean()


if __name__ == "__main__":
    asyncio.run(main())
